#!/bin/bash

# ==============================================================================
# Advanced CPU Affinity Binder for Java Processes
# ==============================================================================
#
# Description:
# This script intelligently binds specified Java processes to CPU cores based on
# the system's hardware topology (Socket, NUMA Node, L3 Cache/CCX).
# It runs continuously, detects new and terminated processes, and rebalances
# CPU affinities automatically to ensure optimal performance.
#
# Features:
#   1. Auto-detects CPU topology (Socket, NUMA, L3 Cache groups).
#   2. Binds processes with a sophisticated priority:
#      - Priority 1: Same L3 Cache (approximates a CCX).
#      - Priority 2: Same NUMA Node.
#      - Priority 3: Same Socket.
#   3. Dynamically manages a pool of available CPUs.
#   4. Monitors for new processes and binds them according to the rules.
#   5. Detects terminated processes, releases their cores, and triggers a
#      re-optimization of all existing bindings.
#
# ==============================================================================

# --- 默认配置 ---
CORES_PER_PROCESS=4
# 用于识别目标Java进程的默认匹配模式。
# 这个模式应该足够精确，以仅匹配到期望的进程。
JAVA_PROCESS_PATTERN="org.apache.spark.executor.CoarseGrainedExecutorBackend"
# 两次检查之间的时间间隔（秒）。
CHECK_INTERVAL=10

# --- 数据结构 ---
# 用于存储系统拓扑和状态的关联数组。
declare -A CPU_TOPOLOGY    # CPU_TOPOLOGY[cpu]="node:socket:l3"
declare -A L3_CORES        # L3_CORES[l3_id]="cpu1 cpu2 ..."
declare -A NODE_CORES      # NODE_CORES[node_id]="cpu1 cpu2 ..."
declare -A SOCKET_CORES    # SOCKET_CORES[socket_id]="cpu1 cpu2 ..."
declare -A BOUND_PROCESSES # BOUND_PROCESSES[pid]="cpu1,cpu2,..."
declare -A USED_CPUS       # USED_CPUS[cpu]=pid

# --- 函数定义 ---

# 打印脚本用法信息。
usage() {
    echo "用法: $0 [-c cores] [-p pattern] [-i interval]"
    echo "  -c, --cores       分配给每个进程的CPU核心数 (默认: $CORES_PER_PROCESS)。"
    echo "  -p, --pattern     用于识别目标Java进程的模式 (例如类名) (默认: '$JAVA_PROCESS_PATTERN')。"
    echo "  -i, --interval    检查进程变化的间隔时间（秒）(默认: $CHECK_INTERVAL)。"
    echo "  -h, --help        显示此帮助信息。"
    exit 1
}

# 解析命令行参数。
parse_args() {
    while [[ "$1" != "" ]]; do
        case $1 in
            -c | --cores )      shift
                                CORES_PER_PROCESS=$1
                                ;;
            -p | --pattern )    shift
                                JAVA_PROCESS_PATTERN=$1
                                ;;
            -i | --interval )   shift
                                CHECK_INTERVAL=$1
                                ;;
            -h | --help )       usage
                                ;;
            * )                 usage
        esac
        shift
    done
}

# 检查所需的命令行工具。
check_dependencies() {
    local missing_deps=0
    for cmd in lscpu pgrep hgvmctl numactl; do
        if ! command -v "$cmd" &> /dev/null; then
            echo "错误：找不到必需的命令 '$cmd'。"
            missing_deps=1
        fi
    done
    if [[ $missing_deps -eq 1 ]]; then
        exit 1
    fi
}

# 从 'lscpu' 和 'numactl' 的输出构建CPU拓扑图。
build_cpu_topology() {
    echo "--- 正在分析系统拓扑 ---"
    # 显示 numactl 和 lscpu 的摘要信息作为参考。
    numactl -H
    echo "---------------------------------"
    lscpu -e
    echo "---------------------------------"
    
    echo "正在构建内部CPU拓扑图..."
    
    # 修正: 强制使用C语言环境，以确保`lscpu`命令输出格式一致，不受系统本地化设置影响。
    export LANG=C

    # 修正: 使用更通用的 'CACHE' 列代替可能不存在的 'L3' 列。
    # CACHE 列的格式通常是 L1d:L1i:L2:L3，我们需要从中解析出L3的ID。
    # 使用进程替换 '< <(...)' 来避免 `while` 循环在子shell中运行，从而保证数组变量在主脚本中可用。
    while IFS=',' read -r cpu node socket cache; do
        # 如果任何值为-，则跳过该行以避免错误（在某些虚拟化环境中可能发生）。
        if [[ -z "$cpu" || -z "$node" || -z "$socket" || -z "$cache" ]]; then
            continue
        fi

        # 从CACHE字符串（例如 "0:0:0:0"）中提取L3缓存ID（结果为 "0"）。
        local l3=${cache##*:}

        CPU_TOPOLOGY[$cpu]="$node:$socket:$l3"
        L3_CORES[$l3]="${L3_CORES[$l3]} $cpu"
        NODE_CORES[$node]="${NODE_CORES[$node]} $cpu"
        SOCKET_CORES[$socket]="${SOCKET_CORES[$socket]} $cpu"
    done < <(lscpu -p=CPU,NODE,SOCKET,CACHE | grep -v '^#')
    
    echo "CPU拓扑分析完成。"
    echo "---------------------------------"
}

# 根据优先级规则查找最佳的可用核心集。
# 返回一个以空格分隔的CPU核心字符串。
find_best_cores() {
    local num_cores=$1
    local best_set=""

    # 优先级 1: 在同一个 L3 缓存组 (CCX) 中寻找 N 个核心。
    for l3_id in "${!L3_CORES[@]}"; do
        local available_in_l3=()
        for cpu in ${L3_CORES[$l3_id]}; do
            [[ -z "${USED_CPUS[$cpu]}" ]] && available_in_l3+=($cpu)
        done
        if (( ${#available_in_l3[@]} >= num_cores )); then
            best_set="${available_in_l3[@]:0:$num_cores}"
            echo "$best_set"
            return 0
        fi
    done

    # 优先级 2: 在同一个 NUMA 节点中寻找 N 个核心。
    for node_id in "${!NODE_CORES[@]}"; do
        local available_in_node=()
        for cpu in ${NODE_CORES[$node_id]}; do
            [[ -z "${USED_CPUS[$cpu]}" ]] && available_in_node+=($cpu)
        done
        if (( ${#available_in_node[@]} >= num_cores )); then
            best_set="${available_in_node[@]:0:$num_cores}"
            echo "$best_set"
            return 0
        fi
    done

    # 优先级 3: 在同一个 Socket 中寻找 N 个核心。
    for socket_id in "${!SOCKET_CORES[@]}"; do
        local available_in_socket=()
        for cpu in ${SOCKET_CORES[$socket_id]}; do
            [[ -z "${USED_CPUS[$cpu]}" ]] && available_in_socket+=($cpu)
        done
        if (( ${#available_in_socket[@]} >= num_cores )); then
            best_set="${available_in_socket[@]:0:$num_cores}"
            echo "$best_set"
            return 0
        fi
    done

    # 最后手段: 在整个系统中寻找任意 N 个可用的核心。
    local available_anywhere=()
    # 对CPU进行数字排序以获得可预测的分配顺序
    local sorted_cpus=($(for cpu in "${!CPU_TOPOLOGY[@]}"; do echo "$cpu"; done | sort -n))
    for cpu in "${sorted_cpus[@]}"; do
        [[ -z "${USED_CPUS[$cpu]}" ]] && available_anywhere+=("$cpu")
    done
    if (( ${#available_anywhere[@]} >= num_cores )); then
        best_set="${available_anywhere[@]:0:$num_cores}"
        echo "$best_set"
        return 0
    fi


    echo "" # 如果找不到合适的CPU块，则返回空。
    return 1
}

# 使用 hgvmctl 将进程绑定到一组核心。
bind_process() {
    local pid=$1
    local core_list_space_separated=$2
    local core_list_comma_separated
    core_list_comma_separated=$(echo "$core_list_space_separated" | tr ' ' ',')

    echo "正在尝试将 PID $pid 绑定到 CPU: $core_list_comma_separated"
    if hgvmctl process bind -c "$core_list_comma_separated" -p "$pid"; then
        BOUND_PROCESSES[$pid]="$core_list_comma_separated"
        for cpu in $core_list_space_separated; do
            USED_CPUS[$cpu]=$pid
        done
        echo "成功绑定 PID $pid。"
    else
        echo "错误: 绑定 PID $pid 失败。该进程可能已终止。"
    fi
}

# 主监控与绑定循环。
main_loop() {
    while true; do
        echo "--- [$(date)]: 开始检查周期 ---"
        local rebalance_needed=false

        # 获取所有匹配模式的PID
        local current_pids=()
        readarray -t current_pids < <(pgrep -f "$JAVA_PROCESS_PATTERN")

        # 1. 检查已终止的进程。
        for pid in "${!BOUND_PROCESSES[@]}"; do
            if ! kill -0 "$pid" 2>/dev/null; then
                echo "进程 $pid 已终止。正在释放其占用的核心。"
                rebalance_needed=true
                # 触发了重新平衡，直接跳出循环，后续会处理所有进程。
                break
            fi
        done

        # 2. 决定本次周期需要处理哪些PID。
        local pids_to_bind=()
        if [[ "$rebalance_needed" == "true" ]]; then
            echo "触发重新平衡：将重新评估所有现有绑定。"
            # 清空所有状态，准备重新绑定所有当前运行的进程。
            BOUND_PROCESSES=()
            USED_CPUS=()
            pids_to_bind=("${current_pids[@]}")
        else
            # 无需重新平衡，仅查找新进程。
            for pid in "${current_pids[@]}"; do
                [[ -z "${BOUND_PROCESSES[$pid]}" ]] && pids_to_bind+=("$pid")
            done
        fi

        # 3. 绑定新进程或重新绑定所有进程。
        if (( ${#pids_to_bind[@]} > 0 )); then
            if [[ "$rebalance_needed" == "true" ]]; then
                 echo "正在重新绑定 ${#pids_to_bind[@]} 个进程..."
            else
                 echo "发现 ${#pids_to_bind[@]} 个新进程需要绑定..."
            fi

            for pid in "${pids_to_bind[@]}"; do
                # 再次检查是否已绑定 (避免竞争条件)
                if [[ -n "${BOUND_PROCESSES[$pid]}" ]]; then
                    continue
                fi

                local core_set
                core_set=$(find_best_cores "$CORES_PER_PROCESS")
                if [[ -n "$core_set" ]]; then
                    bind_process "$pid" "$core_set"
                else
                    echo "警告: 无法为 PID $pid 找到大小为 $CORES_PER_PROCESS 的理想CPU核心块。该进程将保持未绑定状态，直到资源被释放。"
                fi
            done
        else
            echo "未发现新进程，且无需重新平衡。"
        fi

        echo "当前状态:"
        if (( ${#BOUND_PROCESSES[@]} > 0 )); then
             echo "  - 已绑定进程数: ${#BOUND_PROCESSES[@]}"
             # 排序PID以获得一致的输出
             for pid in $(echo "${!BOUND_PROCESSES[@]}" | tr ' ' '\n' | sort -n); do
                printf "    - PID %-10s -> CPUs [%s]\n" "$pid" "${BOUND_PROCESSES[$pid]}"
             done
             echo "  - 已使用核心总数: ${#USED_CPUS[@]}"
        else
            echo "  - 当前脚本未绑定任何进程。"
        fi

        echo "--- 周期结束。休眠 $CHECK_INTERVAL 秒。 ---"
        echo ""
        sleep "$CHECK_INTERVAL"
    done
}


# --- 脚本执行 ---
trap 'echo "\n接收到中断信号，正在退出。"; exit 0' INT TERM

parse_args "$@"
check_dependencies
build_cpu_topology
main_loop

